#!/usr/bin/python
# -*-coding:utf-8-*-
import os
import numpy as np
import pandas as pd

### 底层读取数据的依赖（不提供）
from zbc_factor_lib.base.factors_library_base import NewRQFactorLib as DataReader

# from dir_info import ylabel_data_dir
from dir_info import label_data_dir

from base.KeyTransform import stock_code_map
from base.KeyTransform import trade_date_map
from base.KeyTransform import ci1_code_map
from base.KeyTransform import main as update_map_data_main


def main(db = 'validation'):
    data_reader = DataReader(db=db)

    # TODO - 先更新map数据
    update_map_data_main()

    # TODO - 更新label数据
    label_data = data_reader.read_basic_data_table('processed_new_stock_label_ci1_data')

    size_factor_data = data_reader.read_factor_table(filename='scale_circulate_market_size')

    label_data = label_data.set_index(['date', 'stock_code'])
    size_factor_data = size_factor_data.set_index(['date', 'stock_code'])

    size_factor_data = size_factor_data.rename(columns={'scale_circulate_market_size': 'cap'})

    concat_label_data = pd.concat([label_data,
                                   size_factor_data], axis=1,
                                  join_axes=[label_data.index])

    # TODO - 去掉中信一级行业未分类数据
    concat_label_data = concat_label_data[concat_label_data['ci1_code'] != 'CI005000']

    print(concat_label_data.tail())

    # TODO - 映射
    concat_label_data = concat_label_data.reset_index()

    concat_label_data = trade_date_map(concat_label_data, key='date', reverse=False)
    concat_label_data = stock_code_map(concat_label_data, key='stock_code', reverse=False)
    concat_label_data = ci1_code_map(concat_label_data, key='ci1_code', reverse=False)

    pd.DataFrame(concat_label_data['stock_code'].unique()).to_excel('concat_label_data_stock_code_check.xlsx')

    concat_label_data['cap'] = concat_label_data['cap'].astype(np.float32)

    concat_label_data = concat_label_data.set_index(['date', 'stock_code'])

    print(concat_label_data.tail())

    ## TODO - SAVE
    save_filename = 'label_data'

    concat_label_data.to_hdf(os.path.join(label_data_dir, save_filename + '.h5'),
                             key=save_filename)

if __name__ == "__main__":
    main()


